import numpy as np
from collections import defaultdict
import heapq

class MovieRecommendationKNN:
    def __init__(self, k=3):
        self.k = k
        self.user_ratings = {}
        self.movie_users = defaultdict(set)
    
    def fit(self, ratings_data):
        """
        ratings_data: [(user_id, movie_id, rating), ...]
        """
        for user_id, movie_id, rating in ratings_data:
            if user_id not in self.user_ratings:
                self.user_ratings[user_id] = {}
            self.user_ratings[user_id][movie_id] = rating
            self.movie_users[movie_id].add(user_id)
    
    def cosine_similarity(self, user1, user2):
        """计算用户间的余弦相似度"""
        user1_ratings = self.user_ratings[user1]
        user2_ratings = self.user_ratings[user2]
        
        # 找到共同评分的电影
        common_movies = set(user1_ratings.keys()) & set(user2_ratings.keys())
        
        if len(common_movies) == 0:
            return 0
        
        # 计算余弦相似度
        numerator = sum(user1_ratings[movie] * user2_ratings[movie] 
                       for movie in common_movies)
        
        sum1 = sum(user1_ratings[movie] ** 2 for movie in common_movies)
        sum2 = sum(user2_ratings[movie] ** 2 for movie in common_movies)
        
        denominator = np.sqrt(sum1 * sum2)
        
        if denominator == 0:
            return 0
        
        return numerator / denominator
    
    def find_similar_users(self, target_user):
        """找到最相似的K个用户"""
        similarities = []
        
        for user in self.user_ratings:
            if user != target_user:
                similarity = self.cosine_similarity(target_user, user)
                similarities.append((similarity, user))
        
        # 返回相似度最高的K个用户
        similarities.sort(reverse=True)
        return similarities[:self.k]
    
    def recommend_movies(self, target_user, num_recommendations=5):
        """为目标用户推荐电影"""
        if target_user not in self.user_ratings:
            return []
        
        similar_users = self.find_similar_users(target_user)
        target_user_movies = set(self.user_ratings[target_user].keys())
        
        movie_scores = defaultdict(float)
        similarity_sums = defaultdict(float)
        
        # 基于相似用户的评分进行推荐
        for similarity, similar_user in similar_users:
            for movie, rating in self.user_ratings[similar_user].items():
                if movie not in target_user_movies:  # 推荐未看过的电影
                    movie_scores[movie] += similarity * rating
                    similarity_sums[movie] += similarity
        
        # 计算加权平均评分
        recommendations = []
        for movie in movie_scores:
            if similarity_sums[movie] > 0:
                avg_score = movie_scores[movie] / similarity_sums[movie]
                recommendations.append((avg_score, movie))
        
        # 按评分排序，返回前N个推荐
        recommendations.sort(reverse=True)
        return recommendations[:num_recommendations]
    
    def explain_recommendation(self, target_user, movie):
        """解释推荐理由"""
        similar_users = self.find_similar_users(target_user)
        explanations = []
        
        for similarity, user in similar_users:
            if movie in self.user_ratings[user]:
                rating = self.user_ratings[user][movie]
                explanations.append(f"用户{user}给这部电影评分{rating}分 (相似度: {similarity:.2f})")
        
        return explanations

# 测试电影推荐系统
if __name__ == "__main__":
    # 模拟用户评分数据 (user_id, movie_id, rating)
    ratings_data = [
        # 用户1的评分
        (1, "泰坦尼克号", 5), (1, "阿凡达", 4), (1, "复仇者联盟", 3), (1, "星际穿越", 5),
        # 用户2的评分
        (2, "泰坦尼克号", 4), (2, "阿凡达", 5), (2, "盗梦空间", 5), (2, "星际穿越", 4),
        # 用户3的评分
        (3, "复仇者联盟", 5), (3, "钢铁侠", 4), (3, "蜘蛛侠", 4), (3, "雷神", 3),
        # 用户4的评分
        (4, "泰坦尼克号", 3), (4, "盗梦空间", 5), (4, "星际穿越", 5), (4, "源代码", 4),
        # 用户5的评分
        (5, "复仇者联盟", 4), (5, "钢铁侠", 5), (5, "雷神", 4), (5, "美国队长", 4),
    ]
    
    # 创建推荐系统
    recommender = MovieRecommendationKNN(k=2)
    recommender.fit(ratings_data)
    
    # 为用户1推荐电影
    target_user = 1
    recommendations = recommender.recommend_movies(target_user, num_recommendations=3)
    
    print(f"为用户{target_user}推荐的电影：")
    for score, movie in recommendations:
        print(f"电影: {movie}, 预测评分: {score:.2f}")
        explanations = recommender.explain_recommendation(target_user, movie)
        for explanation in explanations:
            print(f"  - {explanation}")
        print() 
        